import unittest
from kuai.mol.mol import *
from kuai.ff.forcefield import ForceField
from kuai.sim.energy import Energy
from kuai.sim.optimizer import DownhillSimplexOptimizer, PowellOptimizer
from random import random
from math import sqrt

class TestEnergy(Energy):
    def distance(self, mol):
        x = mol.coords[0] - mol.coords[3]
        y = mol.coords[1] - mol.coords[4]
        z = mol.coords[2] - mol.coords[5]
        return sqrt(x*x + y*y + z*z)
        
    def energy(self, mol, ff):
        r = self.distance(mol) - 1
        return r * r

class TestOptimizer(unittest.TestCase):
    TEST_NUMBER = 10
    
    def testDownhill(self):
        ff = ForceField("TEAM")
        energy = TestEnergy()
        for i in range(TestOptimizer.TEST_NUMBER):
            atom1 = Atom("C", XYZ(random(), random(), random()))
            atom2 = Atom("C", XYZ(random(), random(), random()))
            bond = Bond(atom1, atom2)
            mol1 = Molecule([atom1, atom2], [bond])
            opt1 = DownhillSimplexOptimizer()
            opt1.optimze(mol1, ff, energy)
            r = energy.distance(mol1)
            print r
            self.assertAlmostEqual(r, 1, 4)
            
            mol2 = Molecule([atom1, atom2], [bond])
            opt2 = PowellOptimizer()
            opt2.optimze(mol2, ff, energy)
            r = energy.distance(mol2)
            print r
            self.assertAlmostEqual(r, 1, 4)
            
    
